# 引入torch
import torch
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
# pyTorch高阶模块nn
from torch import nn

data = pd.read_csv('dataset/Income.csv')
# print(data)
# 创建pyTorch模型
# 1.预处理数据,转化为框架识别的数据
# 获取数据 特征值
X = data.Education
# 数组拆分为-1自动计算个数，1长度为1的数组
X = X.values.reshape(-1, 1)
# 使用np转换格式
X = X.astype(np.float32)

# 导入数据到torch中成为一个tensor
X = torch.from_numpy(X)

# 处理目标值
Y = torch.from_numpy(data.Income.values.reshape(-1, 1).astype(np.float32))
# 创建模型
# 创建线性层
# output = W@input + b[偏差值] Linear会随机初始化一个w和b来等于output = model（input）
# 新建一个线性模型
model = nn.Linear(1, 1)
# 计算均方误差[损失函数]
loss_fn = nn.MSELoss()
# 优化方法，随机梯度下降算法[SGD],[parameters]会返回随机生成的b和w
# lr = 学习速率
opt = torch.optim.SGD(model.parameters(), lr=0.0001)

# 训练过程，把输入带入到模型中，计算一个结果和真实结果放入损失函数，计算损失
# 使用优化方法，根据损失不断优化w和b
# 全部数据训练一遍叫一次epoch
# 5000次训练
for epoch in range(5000):
    # zip函数是把传入的两个数组，下标一致的元素返回为一个list的类似于Map的对象
    for x, y in zip(X, Y):
        # 使用模型预测
        y_pred = model(x)
        # 根据预测值和真实值计算损失值【误差值】
        loss = loss_fn(y, y_pred)
        # 清空梯度
        opt.zero_grad()
        # 求解权重梯度，反向传播
        loss.backward()
        # 优化变量 w和b
        opt.step()

# 查看权重
print(model.weight)
# 查看偏差
print(model.bias)

# # 绘图 散点图
# 原有数据
plt.scatter(data.Education, data.Income)
# 划线
# 绘出，训练值与预测值的一根线，来和原来的散点比较
plt.plot(X.numpy(), model(X).data.numpy(), c='r')
plt.show()
